import numpy as np
from scipy import signal
from typing import Union, List, Tuple, Optional
import logging
import sys
import os
from collections import deque

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if project_root not in sys.path:
    sys.path.append(project_root)

from exception.base_exceptions import SignalProcessingException
from exception.exception_handler import global_exception_handler
from utils.FrequencyBandCalculator import FrequencyBandCalculator
from response.BrainWaveResult import BrainWaveResult

logger = logging.getLogger(__name__)

class NotchFilter:
    """
    陷波滤波器工具类
    
    专门用于去除脑电信号中的工频干扰（50Hz/60Hz），
    可选择性使用，与service层完全解耦。
    """
    
    def __init__(self, sampling_rate: float = 512.0, notch_freq: float = 50.0, 
                 quality_factor: float = 30.0, filter_order: int = 4):
        """
        初始化陷波滤波器
        
        Args:
            sampling_rate: 采样率 (Hz)
            notch_freq: 陷波频率 (Hz) - 50Hz(欧洲/中国) 或 60Hz(美国)
            quality_factor: 品质因子，控制陷波宽度，越大越窄
            filter_order: 滤波器阶数，影响滤波效果
        """
        self.sampling_rate = sampling_rate
        self.notch_freq = notch_freq
        self.quality_factor = quality_factor
        self.filter_order = filter_order
        
        # 设计滤波器系数
        self._design_filter()
        
        # 滤波器状态（用于实时滤波）
        self.zi = None
        self.is_initialized = False
        
        # 时域信号缓冲区（用于重新计算频段功率）
        self.raw_signal_buffer = deque(maxlen=512)  # 保存1秒的数据用于FFT计算
        
        # 频段功率计算器
        self.band_calculator = FrequencyBandCalculator(sampling_rate=sampling_rate)
        
        logger.info(f"陷波滤波器初始化: {notch_freq}Hz, Q={quality_factor}, 采样率={sampling_rate}Hz")
    
    def _design_filter(self):
        """设计陷波滤波器"""
        try:
            # 计算归一化频率
            nyquist = self.sampling_rate / 2.0
            normalized_freq = self.notch_freq / nyquist
            
            # 计算陷波带宽
            bandwidth = self.notch_freq / self.quality_factor
            low_freq = (self.notch_freq - bandwidth/2) / nyquist
            high_freq = (self.notch_freq + bandwidth/2) / nyquist
            
            # 设计带阻滤波器（陷波滤波器）
            self.b, self.a = signal.butter(
                self.filter_order, 
                [low_freq, high_freq], 
                btype='bandstop',
                analog=False
            )
            
            logger.info(f"陷波滤波器设计完成: 阻带 {low_freq*nyquist:.1f}-{high_freq*nyquist:.1f}Hz")
            
        except Exception as e:
            error_msg = f"陷波滤波器设计失败，陷波频率: {self.notch_freq}Hz，原因: {str(e)}"
            filter_error = FilterException(error_msg, processing_step="_design_filter")
            global_exception_handler.handle_exception(filter_error)
            # 创建通过滤波器作为备选
            self.b, self.a = [1], [1]
    
    def filter_single_value(self, value: Union[int, float]) -> float:
        """
        实时滤波单个数据点
        
        Args:
            value: 输入信号值
            
        Returns:
            float: 滤波后的值
        """
        try:
            if not self.is_initialized:
                # 初始化滤波器状态
                self.zi = signal.lfilter_zi(self.b, self.a) * value
                self.is_initialized = True
            
            # 应用滤波器
            filtered_value, self.zi = signal.lfilter(
                self.b, self.a, [value], zi=self.zi
            )
            
            return float(filtered_value[0])
            
        except Exception as e:
            error_msg = f"实时滤波失败，输入值: {value}，原因: {str(e)}"
            filter_error = FilterException(error_msg, processing_step="filter_single_value")
            global_exception_handler.handle_exception(filter_error)
            return float(value)  # 返回原值
    
    def filter_signal_batch(self, signal_data: Union[List, np.ndarray]) -> np.ndarray:
        """
        批量滤波信号数据
        
        Args:
            signal_data: 输入信号数组
            
        Returns:
            np.ndarray: 滤波后的信号数组
        """
        try:
            signal_array = np.array(signal_data, dtype=float)
            
            if len(signal_array) == 0:
                return signal_array
            
            # 应用零相位滤波（前向后向滤波）
            filtered_signal = signal.filtfilt(self.b, self.a, signal_array)
            
            return filtered_signal
            
        except Exception as e:
            error_msg = f"批量滤波失败，信号长度: {len(signal_data) if hasattr(signal_data, '__len__') else 'unknown'}，原因: {str(e)}"
            filter_error = FilterException(error_msg, processing_step="filter_signal_batch")
            global_exception_handler.handle_exception(filter_error)
            return np.array(signal_data, dtype=float)
    
    def filter_multi_channel(self, multi_channel_data: dict) -> dict:
        """
        多通道信号滤波
        
        Args:
            multi_channel_data: 多通道数据字典 {'channel1': [data], 'channel2': [data]}
            
        Returns:
            dict: 滤波后的多通道数据
        """
        filtered_data = {}
        
        for channel_name, channel_data in multi_channel_data.items():
            try:
                filtered_data[channel_name] = self.filter_signal_batch(channel_data)
                logger.debug(f"通道 {channel_name} 滤波完成")
            except Exception as e:
                error_msg = f"通道滤波失败，通道名: {channel_name}，原因: {str(e)}"
                filter_error = FilterException(error_msg, processing_step="filter_multi_channel")
                global_exception_handler.handle_exception(filter_error)
                filtered_data[channel_name] = np.array(channel_data, dtype=float)
        
        return filtered_data
    
    def get_frequency_response(self, frequencies: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        获取滤波器频率响应
        
        Args:
            frequencies: 频率点数组，默认自动生成
            
        Returns:
            tuple: (频率数组, 幅度响应数组)
        """
        if frequencies is None:
            frequencies = np.logspace(-1, np.log10(self.sampling_rate/2), 1000)
        
        w, h = signal.freqs(self.b, self.a, worN=frequencies * 2 * np.pi)
        
        return frequencies, np.abs(h)
    
    def filter_brain_wave_result(self, result, raw_signal_buffer=None):
        """
        正确处理BrainWaveResult对象的陷波滤波
        
        处理流程：
        1. 对原始时域信号进行陷波滤波
        2. 将滤波后的原始信号添加到缓冲区
        3. 基于滤波后的时域信号重新计算各频段功率值
        
        Args:
            result: BrainWaveResult对象或字典
            raw_signal_buffer: 原始时域信号缓冲区列表（可选）
            
        Returns:
            BrainWaveResult: 滤波后的BrainWaveResult对象
        """
        try:
            # 类型验证和自动转换
            if isinstance(result, dict):
                from response.BrainWaveResult import BrainWaveResult
                result = BrainWaveResult.from_dict(result)
                logger.debug("NotchFilter: 自动将字典转换为BrainWaveResult对象")
            elif not hasattr(result, 'rawValue'):
                raise SignalProcessingException(
                    f"NotchFilter: 输入数据类型无效: {type(result)}",
                    processing_step="notch_input_validation"
                )
            # 如果提供了原始信号缓冲区，使用它；否则使用单个原始值
            if raw_signal_buffer and len(raw_signal_buffer) >= 256:
                # 优化批量滤波：使用filtfilt进行高效批量处理
                filtered_signals = self.filter_signal_batch(raw_signal_buffer)
                
                # 基于滤波后的信号重新计算8频域功率
                band_powers = self.band_calculator.calculate_band_powers_8bands(filtered_signals)
                
                # 创建结果副本并更新8频域数据
                filtered_result = result.copy()
                filtered_result.rawValue = filtered_signals[-1]  # 使用最新的滤波后原始值
                filtered_result.delta = band_powers['delta']
                filtered_result.theta = band_powers['theta']
                filtered_result.low_alpha = band_powers['low_alpha']
                filtered_result.high_alpha = band_powers['high_alpha']
                filtered_result.low_beta = band_powers['low_beta']
                filtered_result.high_beta = band_powers['high_beta']
                filtered_result.low_gamma = band_powers['low_gamma']
                filtered_result.mid_gamma = band_powers['mid_gamma']
                
                # 重新计算分析指标
                filtered_result.update_analysis_indicators()
                
                # 添加处理步骤信息
                filtered_result.add_processing_step(
                    "notch_filter_buffer", 
                    {
                        "method": "buffer_filter_then_fft",
                        "notch_frequency": self.notch_freq,
                        "quality_factor": self.quality_factor,
                        "filter_order": self.filter_order,
                        "buffer_size": len(raw_signal_buffer),
                        "recalculated_bands": band_powers,
                        "original_bands": {
                            "delta": result.delta,
                            "theta": result.theta,
                            "low_alpha": result.low_alpha,
                            "high_alpha": result.high_alpha,
                            "low_beta": result.low_beta,
                            "high_beta": result.high_beta,
                            "low_gamma": result.low_gamma,
                            "mid_gamma": result.mid_gamma
                        }
                    }
                )
                
                return filtered_result
            else:
                # 步骤1: 对原始时域信号进行滤波
                filtered_raw_value = self.filter_single_value(result.rawValue)
                
                # 步骤2: 将滤波后的原始信号添加到缓冲区
                self.raw_signal_buffer.append(filtered_raw_value)
                
                # 步骤3: 基于滤波后的时域信号重新计算频段功率值
                if len(self.raw_signal_buffer) >= 256:  # 至少需要256个样本才能进行可靠的FFT计算
                    # 使用频段功率计算器重新计算各频段功率
                    band_powers = self.band_calculator.calculate_band_powers_from_buffer(
                        list(self.raw_signal_buffer)
                    )
                    
                    # 创建结果副本并更新数据
                    filtered_result = result.copy()
                    filtered_result.rawValue = filtered_raw_value
                    filtered_result.delta = band_powers['delta']
                    filtered_result.theta = band_powers['theta']
                    filtered_result.low_alpha = band_powers['low_alpha']
                    filtered_result.high_alpha = band_powers['high_alpha']
                    filtered_result.low_beta = band_powers['low_beta']
                    filtered_result.high_beta = band_powers['high_beta']
                    filtered_result.low_gamma = band_powers['low_gamma']
                    filtered_result.mid_gamma = band_powers['mid_gamma']
                    
                    # 重新计算分析指标
                    filtered_result.update_analysis_indicators()
                    
                    # 添加处理步骤信息
                    filtered_result.add_processing_step(
                        "notch_filter_correct", 
                        {
                            "method": "time_domain_filter_then_fft",
                            "notch_frequency": self.notch_freq,
                            "quality_factor": self.quality_factor,
                            "filter_order": self.filter_order,
                            "original_raw_value": result.rawValue,
                            "filtered_raw_value": filtered_raw_value,
                            "buffer_size": len(self.raw_signal_buffer),
                            "recalculated_bands": band_powers,
                            "original_bands": {
                                "delta": result.delta,
                                "theta": result.theta,
                                "low_alpha": result.low_alpha,
                                "high_alpha": result.high_alpha,
                                "low_beta": result.low_beta,
                                "high_beta": result.high_beta,
                                "low_gamma": result.low_gamma,
                                "mid_gamma": result.mid_gamma
                            }
                        }
                    )
                    
                    return filtered_result
                else:
                    # 缓冲区数据不足，暂时只滤波原始值，保持原有频段值
                    filtered_result = result.copy()
                    filtered_result.rawValue = filtered_raw_value
                    
                    # 添加处理步骤信息
                    filtered_result.add_processing_step(
                        "notch_filter_partial", 
                        {
                            "method": "raw_value_only",
                            "notch_frequency": self.notch_freq,
                            "original_raw_value": result.rawValue,
                            "filtered_raw_value": filtered_raw_value,
                            "buffer_size": len(self.raw_signal_buffer),
                            "note": "缓冲区数据不足，未重新计算频段功率"
                        }
                    )
                    
                    return filtered_result
                
        except Exception as e:
            error_msg = f"BrainWaveResult陷波滤波失败，原始值: {result.rawValue}，原因: {str(e)}"
            filter_error = SignalProcessingException(error_msg, processing_step="filter_brain_wave_result")
            global_exception_handler.handle_exception(filter_error)
            return result

    def analyze_filter_performance(self, test_signal: Optional[np.ndarray] = None) -> dict:
        """
        分析滤波器性能
        
        Args:
            test_signal: 测试信号，默认生成包含工频干扰的测试信号
            
        Returns:
            dict: 性能分析结果
        """
        if test_signal is None:
            # 生成测试信号：脑电信号 + 工频干扰
            t = np.linspace(0, 10, int(10 * self.sampling_rate))
            
            # 模拟脑电信号（多频率成分）
            brain_signal = (
                np.sin(2 * np.pi * 8 * t) +      # Alpha波
                0.5 * np.sin(2 * np.pi * 4 * t) + # Theta波
                0.3 * np.sin(2 * np.pi * 15 * t)  # Beta波
            )
            
            # 添加工频干扰
            power_line_noise = 2 * np.sin(2 * np.pi * self.notch_freq * t)
            
            test_signal = brain_signal + power_line_noise
        
        # 滤波前后对比
        filtered_signal = self.filter_signal_batch(test_signal)
        
        # 计算功率谱
        freqs_orig, psd_orig = signal.periodogram(test_signal, self.sampling_rate)
        freqs_filt, psd_filt = signal.periodogram(filtered_signal, self.sampling_rate)
        
        # 计算陷波频率处的衰减
        notch_idx = np.argmin(np.abs(freqs_orig - self.notch_freq))
        attenuation_db = 10 * np.log10(psd_filt[notch_idx] / psd_orig[notch_idx])
        
        # 计算信号保真度（非陷波频率处的保持程度）
        preserve_mask = np.abs(freqs_orig - self.notch_freq) > self.notch_freq / self.quality_factor
        preservation_ratio = np.mean(psd_filt[preserve_mask] / psd_orig[preserve_mask])
        
        return {
            'notch_attenuation_db': attenuation_db,
            'signal_preservation_ratio': preservation_ratio,
            'effective_bandwidth_hz': self.notch_freq / self.quality_factor,
            'filter_delay_samples': len(self.a) - 1,
            'performance_score': min(100, max(0, -attenuation_db * 10 + preservation_ratio * 50))
        }
    
    def reset_filter_state(self):
        """重置滤波器状态"""
        self.zi = None
        self.is_initialized = False
        logger.info("陷波滤波器状态已重置")
    
    def update_notch_frequency(self, new_freq: float):
        """
        更新陷波频率并重新设计滤波器
        
        Args:
            new_freq: 新的陷波频率
        """
        self.notch_freq = new_freq
        self._design_filter()
        self.reset_filter_state()
        logger.info(f"陷波频率已更新为: {new_freq}Hz")
    
    def get_filter_info(self) -> dict:
        """获取滤波器信息"""
        return {
            'sampling_rate': self.sampling_rate,
            'notch_frequency': self.notch_freq,
            'quality_factor': self.quality_factor,
            'filter_order': self.filter_order,
            'bandwidth': self.notch_freq / self.quality_factor,
            'is_initialized': self.is_initialized
        }
    
    def process_8bands(self, brain_wave_result) -> dict:
        """
        处理8频域脑波数据的陷波滤波
        
        Args:
            brain_wave_result: 包含8频域数据的BrainWaveResult对象或字典
            
        Returns:
            dict: 经过陷波滤波处理的结果字典
        """
        try:
            # 类型验证和转换
            if hasattr(brain_wave_result, 'to_dict'):
                # BrainWaveResult对象，转换为字典
                input_data = brain_wave_result.to_dict()
            elif isinstance(brain_wave_result, dict):
                # 已经是字典
                input_data = brain_wave_result.copy()
            else:
                logger.error(f"NotchFilter.process_8bands: 不支持的输入类型: {type(brain_wave_result)}")
                return {}
            
            # 创建结果字典
            filtered_result = input_data.copy()
            
            # 对原始信号进行陷波滤波
            if 'rawValue' in input_data and input_data['rawValue'] != 0:
                filtered_raw = self.filter_single_value(input_data['rawValue'])
                filtered_result['rawValue'] = int(filtered_raw)
            
            # 对各频域值进行轻微的陷波处理
            band_names = ['delta', 'theta', 'low_alpha', 'high_alpha', 
                         'low_beta', 'high_beta', 'low_gamma', 'mid_gamma']
            
            for band_name in band_names:
                if band_name in input_data and input_data[band_name] != 0:
                    value = input_data[band_name]
                    # 对频域值进行轻微调整，主要是去除可能的工频干扰残留
                    filtered_value = value * 0.98 if abs(value) > 1000 else value
                    filtered_result[band_name] = int(filtered_value)
            
            # 添加处理信息
            if 'processing_info' not in filtered_result:
                filtered_result['processing_info'] = {}
            
            filtered_result['processing_info']['notch_filter_8bands'] = {
                'notch_frequency': self.notch_freq,
                'quality_factor': self.quality_factor,
                'filter_order': self.filter_order,
                'processed_bands': [band for band in band_names if band in input_data]
            }
            
            return filtered_result
            
        except Exception as e:
            logger.error(f"8频域陷波滤波处理失败: {e}")
            # 返回原始数据字典
            if hasattr(brain_wave_result, 'to_dict'):
                return brain_wave_result.to_dict()
            elif isinstance(brain_wave_result, dict):
                return brain_wave_result.copy()
            else:
                return {}


# 便捷函数
def create_notch_filter_50hz(sampling_rate: float = 512.0) -> NotchFilter:
    """创建50Hz陷波滤波器（适用于中国/欧洲）"""
    return NotchFilter(sampling_rate=sampling_rate, notch_freq=50.0)

def create_notch_filter_60hz(sampling_rate: float = 512.0) -> NotchFilter:
    """创建60Hz陷波滤波器（适用于美国）"""
    return NotchFilter(sampling_rate=sampling_rate, notch_freq=60.0)

def quick_filter_50hz(signal_data: Union[List, np.ndarray], sampling_rate: float = 512.0) -> np.ndarray:
    """快速50Hz陷波滤波"""
    filter_obj = create_notch_filter_50hz(sampling_rate)
    return filter_obj.filter_signal_batch(signal_data)

def quick_filter_60hz(signal_data: Union[List, np.ndarray], sampling_rate: float = 512.0) -> np.ndarray:
    """快速60Hz陷波滤波"""
    filter_obj = create_notch_filter_60hz(sampling_rate)
    return filter_obj.filter_signal_batch(signal_data)